
from llama_index.core.node_parser import SentenceSplitter as InnerSentenceSplitter
from llama_index.core.node_parser import SemanticSplitterNodeParser as InnerSemanticSplitterNodeParser
from typing import List

import json
import re

def CheckQuery(query_str:str) -> int:

    match_str_list = [
        "哪些","几类","几种","几个","几种"
    ]

    pattern = re.compile("|".join(match_str_list))
    match = pattern.search(query_str)

    return 0 if match is None else 512


DataSet = [
    "director",
    "emsplus",
    "rcp",
    "umac"
]


def QueryPart(selected_data_set:int):

    input_file_path = f"ext_slice/inputs/output_rag_keyword_{DataSet[selected_data_set]}.json"

    with open(input_file_path,"r",encoding="utf-8") as file:
        data = json.load(file)
    
    for idx in range(len(data)):
        id = data[idx]["id"]
        print(f"id = {id}")
        query_str = data[idx]["query"]
        ext_len_arg = CheckQuery(query_str)

        selected_file_path = list(data[idx]["selected_file_path"])

        for i in range(len(selected_file_path)):

            one = selected_file_path[i]
            file_path = one["file_path"]
            text_slice = one["text_slice"]
            # slice_lenth = len(text_slice)

            with open(file_path, "r", encoding="utf-8") as f:
                content = str(f.read())

            # 字段拓展
            ori_len = len(text_slice)
            index = content.index(text_slice)
            back_Ext = ext_len_arg + ori_len
            lenth = len(content)
            ext_lenth = min(back_Ext,lenth-index)
            ext_text_slice = content[index:index+ext_lenth]

            # 修改字段
            data[idx]["selected_file_path"][i]["text_slice"] = ext_text_slice
            data[idx]["selected_file_path"][i]["ext_len"] = ext_len_arg

    with open(f"ext_slice/outputs/output_dyn_slice_{DataSet[selected_data_set]}.json","w",encoding="utf-8") as file:
        json.dump(data,file,ensure_ascii=False,indent=4)



if __name__ == "__main__":

    for idx in [0,2,3]:
        QueryPart(idx)
        # break

